package com.heeb.djl.boot.rtsp;

import ai.djl.inference.Predictor;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.output.BoundingBox;
import ai.djl.modality.cv.output.DetectedObjects;
import ai.djl.modality.cv.output.Rectangle;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.repository.zoo.ZooModel;
import com.heeb.djl.boot.criteria.ModelCriteria;
import com.heeb.djl.boot.util.RSAUtil;
import lombok.extern.slf4j.Slf4j;
import org.bytedeco.javacv.*;
import org.bytedeco.opencv.opencv_core.Mat;
import org.bytedeco.opencv.opencv_core.Point;
import org.bytedeco.opencv.opencv_core.Rect;
import org.bytedeco.opencv.opencv_core.Scalar;
import java.awt.image.BufferedImage;
import java.util.List;
import java.util.Random;

import static org.bytedeco.opencv.global.opencv_imgproc.*;

/**
 * @Author bjiang
 * @Description //TODO
 * @Date 2022/1/14 15:45
 * @Version 1.0
 */
@Slf4j
public class RtspUtils {
    /**
     * @Author bjiang
     * @Description //TODO 获取rtsp视频流
     * @Date 15:47 2022/1/14
     * @Version 1.0
     * @Param [url]
     * @return void
     */
    public static FFmpegFrameGrabber getRtspByUrl(String url) {
        log.info("获取rtsp视频流开始");
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(url);
        grabber.setFormat("rtsp");
        String byte2Base64= null;
        try {
            byte2Base64 = RSAUtil.getMsgByRsa("****");
        } catch (Exception e) {
            e.printStackTrace();
        }
        grabber.setOption("rtsp_transport", "tcp");//tcp传输协议
        grabber.setOption("appkey", "****");//海康视频 appkey
        grabber.setOption("secret", byte2Base64);//海康视频 secret
        grabber.setOption("playMode", "0");//初始播放模式：0-预览，1-回放
        grabber.setOption("port", "446");//综合安防管理平台端口，若启用HTTPS协议，默认443
        grabber.setOption("enableHTTPS", "1"); //是否启用HTTPS协议与综合安防管理平台交互，这里总是填1
        grabber.setOption("rtsp_flags", "prefer_tcp");
        grabber.setOption("stimeout", "3000000");
        try {
            grabber.start();
        } catch (FFmpegFrameGrabber.Exception e) {
            e.printStackTrace();
        }
        return grabber;
    }

    /**
     * @Author bjiang
     * @Description //TODO 显示目标监测窗口
     * @Date 16:00 2022/1/14
     * @Version 1.0
     * @Param [grabber]
     * @return void
     */
    public static void showCanvasFrame(String input,String path) throws Exception {
        log.info("获取rtsp视频流开始input={}",input);
        FFmpegFrameGrabber grabber = new FFmpegFrameGrabber(input);
        grabber.setFormat("rtsp");
        String byte2Base64= null;
        try {
            byte2Base64 = RSAUtil.getMsgByRsa("****");
        } catch (Exception e) {
            e.printStackTrace();
        }
        grabber.setOption("rtsp_transport", "tcp");//tcp传输协议
        grabber.setOption("appkey", "****");//海康视频 appkey
        grabber.setOption("secret", byte2Base64);//海康视频 secret
        grabber.setOption("playMode", "0");//初始播放模式：0-预览，1-回放
        grabber.setOption("port", "446");//综合安防管理平台端口，若启用HTTPS协议，默认443
        grabber.setOption("enableHTTPS", "1"); //是否启用HTTPS协议与综合安防管理平台交互，这里总是填1
        grabber.setOption("rtsp_flags", "prefer_tcp");
        grabber.setOption("stimeout", "3000000");
        try {
            grabber.start();
        } catch (FFmpegFrameGrabber.Exception e) {
            e.printStackTrace();
        }
        OpenCVFrameConverter.ToMat converter = new OpenCVFrameConverter.ToMat();
        //CanvasFrame canvas = new CanvasFrame("目标监测"); // 新建一个预览窗口
        /*canvas.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE);
        canvas.setVisible(true);
        canvas.setFocusable(true);
        if (canvas.isAlwaysOnTopSupported()) {
            canvas.setAlwaysOnTop(true);
        }*/

        //FFmpegFrameGrabber recorder = FFmpegFrameGrabber.createDefault("rtsp://localhost:10008/1");
        log.info("开始写入视频流");
        Random random = new Random();
        int url=random.nextInt(1000)  + 1;
        FFmpegFrameRecorder recorder = new FFmpegFrameRecorder("rtsp://19.34.201.228:554/"+url,grabber.getImageWidth(),grabber.getImageHeight(),1);
        recorder.start();
        recorder.setVideoCodec(grabber.getVideoCodec());
        recorder.setFormat(grabber.getFormat());
        recorder.setSampleRate(4410);
        recorder.setFrameRate(grabber.getFrameRate());
        recorder.setVideoBitrate(grabber.getVideoBitrate());
        recorder.setAudioCodec(grabber.getVideoCodec());
        Frame frame = null;
        Criteria<Image, DetectedObjects> criteria =new ModelCriteria().getCriteria(path);
        try (ZooModel model = ModelZoo.loadModel(criteria);
             Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
            int i=0;
            List<DetectedObjects.DetectedObject> items=null;
            //for (;canvas.isVisible() && (frame = grabber.grabImage()) != null; ) {
            for (;(frame = grabber.grabImage()) != null; ) {
                Mat img = converter.convert(frame);
                //drawObject(predictor,img);

                BufferedImage buffImg = mat2BufferedImage(img);
                Image image = ImageFactory.getInstance().fromImage(buffImg);
                int imageWidth = image.getWidth();
                int imageHeight = image.getHeight();
                log.info("imageWidth={},imageHeight={}",imageWidth,imageHeight);
                DetectedObjects detections = predictor.predict(image);
                items = detections.items();
                // 遍历检测物
                for (DetectedObjects.DetectedObject item : items) {
                    BoundingBox box = item.getBoundingBox();
                    Rectangle rectangle = box.getBounds();
                    int x = (int) (rectangle.getX() * imageWidth);
                    int y = (int) (rectangle.getY() * imageHeight);
                    Rect face =
                            new Rect(
                                    x,
                                    y,
                                    (int) (rectangle.getWidth() * imageWidth),
                                    (int) (rectangle.getHeight() * imageHeight));
                    // 绘制矩形区域
                    rectangle(img, face, new Scalar(0, 0, 255, 1));

                    int pos_x = Math.max(face.tl().x() - 10, 0);
                    int pos_y = Math.max(face.tl().y() - 10, 0);
                    String className=item.getClassName()+" " + item.getProbability();
                    // 绘制检测物与可能性
                    putText(
                            img,
                            className,
                            new Point(pos_x, pos_y),
                            FONT_HERSHEY_COMPLEX,
                            1.0,
                            new Scalar(0, 0, 255, 1.0));
                }
                //显示视频
                //canvas.showImage(frame);
                recorder.record(frame);
            }
        }
        grabber.stop();
        grabber.release();
        recorder.stop();
        recorder.release();
    }
    /**
     * @Author bjiang
     * @Description //TODO 绘制检测信息
     * @Date 8:57 2022/1/17
     * @Version 1.0
     * @Param [predictor, img]
     * @return void
     */
    public static void drawObject(Predictor<Image, DetectedObjects> predictor,Mat img) throws Exception {

        BufferedImage buffImg = mat2BufferedImage(img);
        Image image = ImageFactory.getInstance().fromImage(buffImg);
        int imageWidth = image.getWidth();
        int imageHeight = image.getHeight();
        log.info("imageWidth={},imageHeight={}",imageWidth,imageHeight);
        DetectedObjects detections = predictor.predict(image);
        List<DetectedObjects.DetectedObject> items = detections.items();
        // 遍历检测物
        for (DetectedObjects.DetectedObject item : items) {
            BoundingBox box = item.getBoundingBox();
            Rectangle rectangle = box.getBounds();
            int x = (int) (rectangle.getX() * imageWidth);
            int y = (int) (rectangle.getY() * imageHeight);
            Rect face =
                    new Rect(
                            x,
                            y,
                            (int) (rectangle.getWidth() * imageWidth),
                            (int) (rectangle.getHeight() * imageHeight));
            // 绘制矩形区域
            rectangle(img, face, new Scalar(0, 0, 255, 1));

            int pos_x = Math.max(face.tl().x() - 10, 0);
            int pos_y = Math.max(face.tl().y() - 10, 0);
            String className=item.getClassName()+" " + item.getProbability();
            // 绘制检测物与可能性
            putText(
                    img,
                    className,
                    new Point(pos_x, pos_y),
                    FONT_HERSHEY_COMPLEX,
                    1.0,
                    new Scalar(0, 0, 255, 2.0));
        }
    }

    /**
     * 将mat转BufferedImage
     *
     */
    public static BufferedImage mat2BufferedImage(Mat matrix) {
        int cols = matrix.cols();
        int rows = matrix.rows();
        int elemSize = (int) matrix.elemSize();
        byte[] data = new byte[cols * rows * elemSize];

        matrix.data().get(data);

        int type = 0;
        switch (matrix.channels()) {
            case 1:
                type = BufferedImage.TYPE_BYTE_GRAY;
                break;
            case 3:
                type = BufferedImage.TYPE_3BYTE_BGR;
                byte b;
                for (int i = 0; i < data.length; i = i + 3) {
                    b = data[i];
                    data[i] = data[i + 2];
                    data[i + 2] = b;
                }
                break;
            default:
                return null;
        }
        BufferedImage image = new BufferedImage(cols, rows, type);
        image.getRaster().setDataElements(0, 0, cols, rows, data);
        return image;
    }
}
